package cn.net.xyan.facerecognition.service

import cn.net.xyan.facerecognition.dao.UserDao
import cn.net.xyan.facerecognition.utils.convert2Mat
import org.bytedeco.javacpp.DoublePointer
import org.bytedeco.javacpp.IntPointer
import org.bytedeco.javacpp.opencv_core.*
import org.bytedeco.javacpp.opencv_face.createLBPHFaceRecognizer
import org.springframework.beans.factory.annotation.Autowired
import java.awt.image.BufferedImage
import java.io.ByteArrayInputStream
import java.nio.IntBuffer
import javax.annotation.PostConstruct
import javax.imageio.ImageIO


interface FaceRecognitionService {
    fun predict(image: BufferedImage): PredictResult
    fun update(image: BufferedImage, label: Int)
}

data class PredictResult(var label: Int?, var confidence: Double?)

class FaceTrainHelper( size: Int) {
    val labels = Mat(size, 1, CV_32SC1)
    lateinit var  labelsBuf: IntBuffer
    val images = MatVector(size.toLong())

    var counter: Int = 0

    init {
        labelsBuf = labels.createBuffer()
    }

    fun addFace(image: BufferedImage, label: Int) {
        images.put(counter.toLong(), image.convert2Mat())
        labelsBuf.put(counter, label)
        counter += 1
    }

    fun close(){
        labels.close()
    }
}


open class OpenCVLBPHFaceRecognizerImpl : FaceRecognitionService {

    @Autowired
    lateinit var userDao: UserDao

    var faceRecognizer = createLBPHFaceRecognizer()

    @PostConstruct
    fun setup() {
        val faces = userDao.findAll().flatMap { it.faces }

        if (faces.size  == 0) return

        val trainHelper = FaceTrainHelper(faces.size)

        faces.forEach {
            face ->
            ByteArrayInputStream(face.imageData).use {
                val image = ImageIO.read(it)

                trainHelper.addFace(image, face.label!!)
            }
        }

        faceRecognizer.train(trainHelper.images, trainHelper.labels)
    }

    override fun update(image: BufferedImage, label: Int) {
        val trainHelper = FaceTrainHelper(1)
        trainHelper.addFace(image, label)
        faceRecognizer.update(trainHelper.images, trainHelper.labels)
    }

    override fun predict(image: BufferedImage): PredictResult {
        val label = IntPointer(1)
        val confidence = DoublePointer(1)
        faceRecognizer.predict(image.convert2Mat(), label, confidence)
        //val predictedLabel = label.get(0)
        return PredictResult(label.get(0), confidence.get(0))
    }
}